import argparse
import torch
import os
import sys

if os.getcwd() not in sys.path:
    sys.path.append(os.getcwd())

from utils.utils import setup_seed
from utils.utils import get_source_list
from utils.utils import list_to_files

import kmeans.kmeans_utils as K
import tqdm
import numpy as np
import yaml
from random import shuffle

SEED = 1234


def main(args):
    setup_seed(SEED)
    embedding_path: list = get_source_list(args.emb_scp)
    shuffle(embedding_path)
    with open(args.config, "r") as file:
        config: dict = yaml.safe_load(file)
    name = (
        os.path.basename(args.config).replace(".yaml", "")
        if args.name is None
        else args.name
    )
    os.makedirs(args.ckpt_dir, exist_ok=True)

    kmeans_model = K.fetch_kmeans_model(
        n_clusters=config["n_clusters"],
        init=config["init"],
        max_iter=config["max_iter"],
        batch_size=config["batch_size"],
        tol=config["tol"],
        max_no_improvement=config["max_no_improvement"],
        n_init=config["n_init"],
        reassignment_ratio=config["reassignment_ratio"],
        random_state=SEED,
        checkpoint_path=os.path.join(
            args.ckpt_dir, f"kmeans-cluster-{config['n_clusters']}-{name}.pt"
        ),
    )
    K.train(
        kmeans_model,
        embedding_path,
        args.ckpt_dir,
        name,
        kmeans_batch_size=config["batch_size"],
    )
    pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--emb_scp",
        type=str,
        required=True,
        help="""
            The file has a list of source paths to wavlm embedding, where each path is on each line. It looks like:
            ...
            >>> path/to/wav/1.npy
            >>> path/to/wav/2.npy
            ...
            """,
    )
    parser.add_argument("--config", type=str, required=True, help="The config path.")
    parser.add_argument(
        "--ckpt_dir", type=str, default = None, help="The checkpoint directory"
    )
    parser.add_argument(
        "--name",
        default=None,
        help="The name of the experiment. if None, use the basename of config as the experiment name.",
    )
    # parser.add_argument("--ckpt", type = str, default = None, help = "The checkpoint")
    args = parser.parse_args()
    if args.ckpt_dir is None:
        args.ckpt_dir = os.path.join(
            os.path.dirname(os.path.dirname(args.config)),
            "ckpt",
            os.path.basename(args.config).replace(".yaml", ""),
        )
    main(args)
    pass
